#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文档检查和图片提取工具
支持DOC和DOCX格式的文档
生成的文件将保存到extract目录
"""

from docx import Document
import os
import zipfile
import shutil
import win32com.client
import pythoncom

# 确保上级extract目录存在
os.makedirs("..\\extract", exist_ok=True)

def read_docx(file_path):
    doc = Document(file_path)
    text = []
    for para in doc.paragraphs:
        text.append(para.text)
    return '\n'.join(text)

def extract_images_from_docx(file_path, output_dir="..\\extract"):
    """
    从Word文档中提取所有图片
    
    Args:
        file_path: Word文档路径
        output_dir: 图片输出目录
    
    Returns:
        list: 提取的图片文件路径列表
    """
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 临时解压目录
    temp_dir = "temp_docx_extract"
    
    try:
        # 解压docx文件（docx本质上是zip文件）
        with zipfile.ZipFile(file_path, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # 图片通常存储在word/media目录中
        media_dir = os.path.join(temp_dir, "word", "media")
        
        extracted_images = []
        
        if os.path.exists(media_dir):
            # 获取所有图片文件
            image_files = [f for f in os.listdir(media_dir) 
                          if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.emf', '.wmf'))]
            
            # 复制图片到输出目录
            for img_file in image_files:
                src_path = os.path.join(media_dir, img_file)
                dst_path = os.path.join(output_dir, img_file)
                shutil.copy2(src_path, dst_path)
                extracted_images.append(dst_path)
                print(f"提取图片: {img_file}")
        
        # 获取文档中的图片信息（通过python-docx）
        doc = Document(file_path)
        
        # 查找文档中的图片
        image_count = 0
        for rel in doc.part.rels.values():
            if "image" in rel.reltype:
                image_count += 1
                print(f"发现图片关系: {rel.reltype}")
        
        print(f"文档中共发现 {len(extracted_images)} 张图片")
        print(f"文档关系中共发现 {image_count} 个图片引用")
        
        return extracted_images
        
    except Exception as e:
        print(f"提取图片时出错: {e}")
        return []
    
    finally:
        # 清理临时目录
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)

def get_document_images_info(file_path):
    """
    获取文档中图片的详细信息
    """
    doc = Document(file_path)
    
    print("=== 文档图片信息 ===")
    
    # 检查段落中的图片
    for i, paragraph in enumerate(doc.paragraphs):
        for run in paragraph.runs:
            if run._element.xpath('.//pic:pic'):
                print(f"段落 {i+1} 中发现图片")
    
    # 检查表格中的图片
    for table_idx, table in enumerate(doc.tables):
        for row_idx, row in enumerate(table.rows):
            for cell_idx, cell in enumerate(row.cells):
                for para in cell.paragraphs:
                    for run in para.runs:
                        if run._element.xpath('.//pic:pic'):
                            print(f"表格 {table_idx+1}, 行 {row_idx+1}, 列 {cell_idx+1} 中发现图片")

if __name__ == '__main__':
    file_path = '../files/2_Network Intrusion Detection with DAECNN Model Based on imbalanced dataset.docx'
    
    # 检查文件是否存在
    if not os.path.exists(file_path):
        print(f"测试文件不存在: {file_path}")
        print("请将测试文档放入files目录或指定正确的文件路径")
        file_path = input("请输入文档路径: ").strip()
        if not os.path.exists(file_path):
            print("文件不存在，程序退出")
            exit(1)
    
    # 提取文本
    ref_text = read_docx(file_path)
    print("=== 文档文本内容 ===")
    print(ref_text[:500] + "..." if len(ref_text) > 500 else ref_text)
    print("\n")
    
    # 提取图片到上级extract目录
    images = extract_images_from_docx(file_path, "..\\extract")
    
    # 获取图片信息
    get_document_images_info(file_path)
    
    if images:
        print(f"\n成功提取 {len(images)} 张图片到 '../extract' 目录:")
        for img in images:
            print(f"  - {os.path.basename(img)}")
    else:
        print("\n未发现图片或提取失败")